fix(loss): min/max packed extrema instead of summing them - #3901
Open
tianyi-zhang-02 wants to merge 3 commits into
Open
fix(loss): min/max packed extrema instead of summing them#3901tianyi-zhang-02 wants to merge 3 commits into
tianyi-zhang-02 wants to merge 3 commits into
Conversation
SequencePackingLossWrapper folds per-sequence metric dicts into one. It special-cases extrema through a hardcoded allowlist naming only the four probs_ratio keys, so MseValueLossFn's values_min/values_max fall through to '+=' and are summed. Three packed sequences spanning -3..9 report values_min=4.0 -- a positive number for a critic whose predictions go negative -- and values_max=15.0. The error grows with packing density, so it is not a stable offset a reader could correct for. The loss function already says what these are: 'Min/max are per-MB; ppo.py takes min/max across MBs.' Five other sites apply the '_min'/'_max' suffix rule to this very dict, and one of them -- megatron_value_worker.py:611 -- is this wrapper's own direct consumer, skipping the divide because it is an extremum while the wrapper upstream has already summed it. Use the same rule here. That alone is not enough. MseValueLossFn returns 0.0 for a fully-masked sequence where ClippedPGLossFn returns +/-inf, and 0.0 is a plausible value that wins the min against an all-positive critic: the one-line version reports 0.0 where the truth is 3.0. sample_mask is loss_multiplier, which overlong_filtering zeroes per sample, and under packing one filtered sample in a pack is enough. So the sentinel moves to +/-inf and ppo.py skips it, matching what that file already does for probs_ratio at :1774 and :2757. Metrics only -- packed and unpacked losses are bit-identical. Reachable on ppo-qwen2.5-1.5b-gsm8k-1n8g-megatron-valuetp2sp-pp2cp2-pack, which nightly.txt runs. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
tianyi-zhang-02
force-pushed
the
fix/packed-metric-extrema
branch
from
September 12, 2026 18:25
43f85ae to
4a5af89
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes packed min/max diagnostics that were being summed as additive metrics.
SequencePackingLossWrapperand the worker aggregators recognize extrema by the explicit_min/_maxsuffix contract. Names that merely contain those substrings in the middle remain additive. Current main's registered MOPD exception,opd_full_decomposition_error, keeps its max reduction even though it does not use the suffix. Fully masked value samples retain the existing sentinel behavior, and an all-sentinel critic step maps to-1.0like policy metrics.Losses and gradients are unchanged.
Validation
Current head
4a5af89017bd2bf51a9fa4cd32ba81ec59de7f20, rebased onto upstreammainat90a2a212d503455d8590be5c8de3cb989d3425b0.git diff --check: passedAn earlier pre-refresh head also passed the gated Megatron split-parity test on two real H100s. It ran two optimizer steps through freshly initialized sync and split policies and compared the loss curve, grad norm, and every reduced per-microbatch metric.
Environment: 2× NVIDIA H100 80 GB,
nvcr.io/nvidia/nemo-rl:v0.7.0, Python 3.13, PyTorch 2.11.0+cu130, CUDA 13, andNCCL_NVLS_ENABLE=0. The current-head rebase and MOPD regression were verified locally; the 2-GPU parity test was not rerun after that refresh.